from Models.models.involution_cuda import involution, _involution_cuda
import torch.nn.functional as F

class MyInvolution(involution):
    """
        继承至 involution 类，主要用于添加残差连接
    """
    def __init__(self, dim, kernel_size, stride, use_res=True, use_softmax=True):
        super().__init__(dim, kernel_size, stride)
        self.use_res = use_res
        self.use_softmax = use_softmax

    def forward(self, x):
        if self.use_res:
            res = x
        weight = self.conv2(self.conv1(x if self.stride == 1 else self.avgpool(x)))
        b, c, h, w = weight.shape
        if self.use_softmax:
            weight = weight.view(b, self.groups, self.kernel_size*self.kernel_size, h, w)
            weight = F.softmax(weight, dim=2)
        weight = weight.view(b, self.groups, self.kernel_size, self.kernel_size, h, w)
        out = _involution_cuda(x, weight, stride=self.stride, padding=(self.kernel_size-1)//2)
        if self.use_res:
            out += res
        return out